#!/usr/bin/env python3
"""
RAG 系统使用示例
演示如何使用检索工具进行文档检索
"""

import os
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.retrieval.retrieval_tool import RetrievalTool
from src.data_processing.vector_store_builder import VectorStoreBuilder
from src.models.llm_config import LLMConfig
from src.utils.logger import setup_logger

logger = setup_logger(__name__)


def main():
    """主函数 - 演示 RAG 系统的基本使用"""
    
    try:
        # 1. 加载配置
        logger.info("加载配置...")
        config = LLMConfig()
        logger.info(f"配置加载成功: {config}")
        
        # 2. 获取向量索引
        logger.info("获取向量索引...")
        builder = VectorStoreBuilder(config)
        index = builder.get_index()
        logger.info("向量索引获取成功")
        
        # 3. 创建检索工具
        logger.info("创建检索工具...")
        retrieval_tool = RetrievalTool(index, config)
        logger.info("检索工具创建成功")
        
        # 4. 执行检索示例
        test_queries = [
            "成都旅游景点推荐",
            "贵阳美食介绍", 
            "交通出行信息",
            "住宿推荐"
        ]
        
        for query in test_queries:
            logger.info(f"\n{'='*50}")
            logger.info(f"查询: {query}")
            logger.info(f"{'='*50}")
            
            # 执行混合检索
            results = retrieval_tool.search(
                query=query,
                search_type="hybrid",
                top_k=3,
                deduplicate=True,
                rerank=True,
                min_score=0.3
            )
            
            if results:
                logger.info(f"找到 {len(results)} 个相关结果:")
                for i, result in enumerate(results, 1):
                    logger.info(f"\n结果 {i}:")
                    logger.info(f"  分数: {result['score']:.3f}")
                    logger.info(f"  内容: {result['text'][:150]}...")
                    logger.info(f"  元数据: {result['metadata']}")
            else:
                logger.warning("未找到相关结果")
        
        # 5. 演示获取上下文
        logger.info(f"\n{'='*50}")
        logger.info("演示获取上下文功能")
        logger.info(f"{'='*50}")
        
        context = retrieval_tool.get_context(
            query="成都旅游攻略",
            max_length=1000
        )
        
        if context:
            logger.info(f"上下文长度: {len(context)} 字符")
            logger.info(f"上下文内容:\n{context}")
        else:
            logger.warning("未获取到上下文")
        
        # 6. 演示元数据过滤
        logger.info(f"\n{'='*50}")
        logger.info("演示元数据过滤功能")
        logger.info(f"{'='*50}")
        
        filtered_results = retrieval_tool.search_with_metadata(
            query="旅游景点",
            metadata_filter={"source": "马蜂窝成都.pdf"}
        )
        
        logger.info(f"过滤后结果数量: {len(filtered_results)}")
        for result in filtered_results:
            logger.info(f"  来源: {result['metadata'].get('source', 'Unknown')}")
            logger.info(f"  内容: {result['text'][:100]}...")
        
        # 7. 测试答案生成功能
        logger.info(f"\n{'='*50}")
        logger.info("测试答案生成功能")
        logger.info(f"{'='*50}")
        
        demo_questions = [
            "成都有什么必去的旅游景点？",
            "贵阳有什么特色美食？",
            "成都哪里可以喝茶？"
        ]
        
        for question in demo_questions:
            logger.info(f"\n问题: {question}")
            logger.info("-" * 40)
            
            answer = retrieval_tool.generate_answer(question)
            logger.info(f"回答: {answer}")
        
        # 8. 获取统计信息
        logger.info(f"\n{'='*50}")
        logger.info("系统统计信息")
        logger.info(f"{'='*50}")
        
        stats = retrieval_tool.get_statistics()
        for key, value in stats.items():
            logger.info(f"{key}: {value}")
        
        logger.info("\nRAG 系统演示完成！")
        
    except Exception as e:
        logger.error(f"演示过程中发生错误: {e}")
        logger.error("请检查:")
        logger.error("1. 环境变量配置是否正确")
        logger.error("2. Qdrant 服务是否正常运行")
        logger.error("3. 知识库是否已构建")
        raise


if __name__ == "__main__":
    main()
